// main.rs
mod test1;

// #[tokio::main]
// async fn main() -> Result<(), Box<dyn std::error::Error>> {
//     // 运行测试
//     test1::run_performance_test_01().await?;
//     Ok(())
// }
// main.rs

use tokio::time::{Duration, Instant};
use zeromq::{prelude::*, *};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use log::{debug, info};
use tokio::spawn;

static TOTAL_MESSAGES: AtomicU64 = AtomicU64::new(0);
static TOTAL_MESSAGES_LEN: AtomicU64 = AtomicU64::new(0);

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    env_logger::builder()
        .filter_level(log::LevelFilter::Info)
        .init();
    let mut handles = vec![];
    // 移除代理服务器启动代码
    // let _proxy_handle = spawn(async move {
    //     if let Err(e) = start_proxy_server().await {
    //         eprintln!("Proxy error: {}",  e);
    //     }
    // });
    //handles.push(_proxy_handle);
    // 等待服务器启动
    tokio::time::sleep(Duration::from_millis(100)).await;

    // 创建生产者和消费者
    let producer_count = 60;
    let consumer_count = 60;


    // 启动消费者
    for i in 0..consumer_count {
        tokio::time::sleep(Duration::from_millis(100)).await;
        info!("启动消费者 {}", i);
        let handle = spawn(async move {
            if let Err(e) = start_pull_consumer(i).await {
                eprintln!("Consumer {} error: {}", i, e);
            }
        });
        handles.push(handle);
    }

    // 等待消费者连接
    tokio::time::sleep(Duration::from_millis(100)).await;

    // 启动生产者
    for i in 0..producer_count {
        //tokio::time::sleep(Duration::from_millis(100)).await;
        info!("启动生产者 {}", i);
        let handle = spawn(async move {
            if let Err(e) = start_push_producer(i).await {
                eprintln!("Producer {} error: {}", i, e);
            }
        });
        //handles.push(handle);
    }
    info!("启动消费者完成");
    let start_time = Instant::now();
    info!("开始测试...");

    // 等待所有任务完成
    for handle in handles {
        handle.await?;
    }

    // 注意：这里不再等待 proxy_task，因为它是无限运行的。
    // 如果需要优雅关闭，应引入 shutdown 机制。

    let duration = start_time.elapsed();
    let total_messages = TOTAL_MESSAGES.load(Ordering::Relaxed);
    let total_msg_len = TOTAL_MESSAGES_LEN.load(Ordering::Relaxed);

    info!("高性能测试完成:");
    info!("总消息数: {}", total_messages);
    info!("耗时: {:?}", duration);
    info!("吞吐量: {:.2} 消息/秒", total_messages as f64 / duration.as_secs_f64());
    info!("总消息长度: {}", total_msg_len);
    info!("吞吐量: {}字节/秒", total_msg_len as f64 / duration.as_secs_f64());

    Ok(())
}



/// 启动一个推送生产者，向指定端口发送大量消息
///
/// 该函数创建一个ZeroMQ请求套接字，连接到指定端口，并发送100万条消息。
/// 每条消息都等待接收方的回复确认。
///
/// # 参数
/// * `id` - 生产者的唯一标识符，用于计算连接端口号和消息标识
///
/// # 返回值
/// * `Result<(), Box<dyn std::error::Error>>` - 成功时返回Ok(())，失败时返回错误信息
async fn start_push_producer(id: usize) -> Result<(), Box<dyn std::error::Error>> {
    info!("启动推送生产者 {}", id);
    let mut socket = zeromq::ReqSocket::new();
    // 修改连接地址为消费者直接监听的端口 5558
    socket.connect(format!("tcp://127.0.0.1:{}",5558+id).as_str()).await?;

    let message_count = 100000; // 每个生产者发送更多消息

    // 循环发送指定数量的消息，每发送10%数量的消息记录一次日志
    let mut message_size = 0u64;
    for i in 0..message_count {
        if i *50 % message_count == 0 {
            info!("生产者 {} 已发送 {} 条消息,{}", id, i, message_size);
        }
        let payload = format!("Producer-{}-Message-{}", id, i);
        message_size += payload.len() as u64;
        let msg = zeromq::ZmqMessage::from(payload);

        // 发送请求消息
        let response =    socket.send(msg).await;
        if let Err(e) = response {
            eprintln!("生产者 {} 发送错误: {}", id, e);
            break;
        }

        // 等待回复
        let _reply = socket.recv().await?;
        message_size += _reply.len() as u64;
    }

    info!("生产者 {} 完成发送 {} 条消息，{}", id, message_count,message_size);
    TOTAL_MESSAGES_LEN.fetch_add(message_size, Ordering::Relaxed);

    Ok(())
}

async fn start_pull_consumer(id: usize) -> Result<(), Box<dyn std::error::Error>> {
    let mut socket = zeromq::RepSocket::new();
    // 修改监听端口为生产者直接连接的端口 5558
    socket.bind(format!("tcp://127.0.0.1:{}",(5558+id)).as_str()).await?;

    info!("消费者 {} 已连接并开始接收消息", id);

    let mut count = 0u64;
    let start_time = Instant::now();
    // 运行10秒钟测试性能
    while start_time.elapsed() < Duration::from_secs(30) {
        match tokio::time::timeout(Duration::from_millis(100), socket.recv()).await {
            Ok(Ok(msg)) => {
                count += 1;

                // 发送回复
                let reply = zeromq::ZmqMessage::from("OK");
                socket.send(reply).await?;

                // 每10000条消息打印一次进度
                if count % 10000 == 0 {
                    debug!("消费者 {} 已接收 {} 条消息", id, count);
                }
            }
            Ok(Err(e)) => {
                eprintln!("消费者 {} 接收错误: {}", id, e);
                break;
            }
            Err(_) => {
                // 超时，继续尝试接收
                continue;
            }
        }
    }
    TOTAL_MESSAGES.fetch_add(count, Ordering::Relaxed);

    info!("消费者 {} 在 {:?} 内接收了 {} 条消息",
             id, start_time.elapsed(), count);
    Ok(())
}
